import os
import pandas as pd


def generate_csv(root_dir, mode):
    if mode == 'train':
        # videos_path = os.path.join(root_dir, 'Train_files')
        videos_path = os.path.join(root_dir, 'Train_files_MSR')
    elif mode == 'val':
        # videos_path = os.path.join(root_dir, 'Dev_files')
        videos_path = os.path.join(root_dir, 'Dev_files_MSR')
    else:
        # videos_path = os.path.join(root_dir, 'Test_files')
        videos_path = os.path.join(root_dir, 'Test_files_MSR')
    fold_data = []
    for video_name in os.listdir(videos_path):
        video_path = os.path.join(videos_path, video_name)
        for image in os.listdir(video_path):
            video_name_split = video_name.split('_')
            assert len(video_name_split) == 4
            label = 0 if int(video_name_split[-1]) == 1 else 1
            fold_data.append([video_name, image, label])
    columns = ["video", "file", "label"]
    # pd.DataFrame(fold_data, columns=columns).to_csv('../data/data_{}.csv'.format(mode), index=False)
    pd.DataFrame(fold_data, columns=columns).to_csv('../data/data_{}_MSR.csv'.format(mode), index=False)



def main():
    # root_dir = '/data/heyan/Datasets/Face/oulu_256'  # RGB images
    root_dir = '/data2/shaohua/Datasets/Oulu_NPU'   # MSR images

    generate_csv(root_dir, 'train')
    generate_csv(root_dir, 'val')
    generate_csv(root_dir, 'test')


if __name__ == '__main__':
    main()
